import base64
import io

import torch
from PIL import Image
from flask import Flask, request, jsonify
from flask import render_template
from torchvision import transforms

app = Flask(__name__)

# 加载预训练的模型
# 确保 SimpleCNN 已经正确定义或导入
# model = SimpleCNN()  # 确保 SimpleCNN 已经正确定义或导入
# model.load_state_dict(torch.load('full_model.pth'))
model = torch.load('full_model2.pth')


# model.eval()  # 设置模型为评估模式


@app.route('/')
def home():
    return render_template('index.html')


# 将图像转换为张量并进行必要的预处理

# 定义图像预处理
transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.Grayscale(),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])


@app.route('/predict', methods=['POST'])
def predict():
    # 获取前端发送的 base64 图片数据
    data = request.form['image']
    header, encoded = data.split(",", 1)
    image_data = base64.b64decode(encoded)

    # 将 base64 数据转换为 PIL 图像
    image = Image.open(io.BytesIO(image_data))
    # 【原图】保存图片到本地
    save_path = "output_origin.png"  # 保存的文件路径
    image.save(save_path)
    print(f"图片已保存到: {save_path}")

    # 将图像转换为灰度图像
    image = image.convert('L')
    # 【灰阶】保存图片到本地
    save_path = "output.png"  # 保存的文件路径
    image.save(save_path)
    print(f"图片已保存到: {save_path}")

    # 对图像进行预处理
    image_tensor = transform(image).unsqueeze(0)
    # 增加 batch 维度
    # 将图像张量移动到 GPU（如果可用）
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    image_tensor = image_tensor.to(device)

    # 使用模型进行预测
    with torch.no_grad():
        output = model(image_tensor)
        _, predicted = torch.max(output, 1)
        predicted_digit = predicted.item()

    # 返回预测结果
    result = jsonify({'predicted_digit': predicted_digit})
    return result


if __name__ == '__main__':
    app.run(debug=True)
